The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

The mixin currently have a set of very relevant model card metadata. For more customization, it's possible to overwrite the generate_model_card method. This method allows more flexibility and access to model's values/config. It return a ModelCard object. For example:

class Model(nn.Module, PytorchModelHubMixin):
    ...

    def generate_model_card(self) -> ModelCard:
        card = super(self).generate_model_card()
        card.data.custom_key = "custom value"
        return card

I'd prefer not to add too many features to the main way of defining tags/library_name/... to avoid complexity and promote this flexible method instead. WDYT?

The feature is inspired by this PR, as you may see I try to make it possible to pass a custom metrics dictionary and dataset name to be able to generate cards like this.
https://huggingface.co/qubvel-hf/oxford-pet-segmentation

model.save_pretrained('./my_model', metrics={'accuracy': 0.95}, dataset='my_dataset')

You might want to pass arguments not at the time of model class creation, but at the time of model saving, for example when the model is trained. I overridden the generate_model_card() method but I didn't find a way to pass any parameters from the save_pretrained() method. The only way I found is to set them as private attributes in save_pretrained and then remove them.
https://github.com/qubvel/segmentation_models.pytorch/blob/3d6da1d74636873372c265f300862a6a6d01777d/segmentation_models_pytorch/base/hub_mixin.py#L107
Now it looks like a dirty hack :(

@Wauplin thank you for the feedback! I will update the integration, would be cool to include this PR changes, to get rid of private attributes :)

Could you add a small test in test_hub_mixin_pytorch.py with a new class that overwrites generate_model_card + pass data in save_pretrained and check you have the expected result. Thanks in advance!

Done! Please, have a look!

Thanks for the quick review and merge 🤗